from typing import List, Optional
import json
from langchain_core.messages import BaseMessage, message_to_dict, messages_from_dict
from langchain_community.chat_message_histories.redis import RedisChatMessageHistory


class LimitedRedisChatMessageHistory(RedisChatMessageHistory):
    """RedisChatMessageHistory with message count limitation"""

    def __init__(
            self,
            session_id: str,
            url: str = "redis://localhost:6379/0",
            key_prefix: str = "message_store:",
            max_messages: int = 10,
            ttl: Optional[int] = None,
            *args,
            **kwargs
    ):
        super().__init__(
            session_id=session_id,
            url=url,
            key_prefix=key_prefix,
            ttl=ttl,
            *args,
            **kwargs
        )
        self.max_messages = max_messages

    def add_message(self, message: BaseMessage) -> None:
        """Add a message and prune if necessary"""
        super().add_message(message)
        self._prune_messages()

    def _prune_messages(self) -> None:
        """Remove oldest messages if total exceeds max_messages using Redis LTRIM"""
        current_length = self.redis_client.llen(self.key)
        if current_length > self.max_messages:
            keep_from = current_length - self.max_messages
            # self.redis_client.ltrim(self.key, keep_from, -1)
            self.redis_client.ltrim(self.key, 0, -self.max_messages - 1)

    def _deserialize_message(self, message_data: bytes) -> BaseMessage:
        """Deserialize Redis message data to BaseMessage"""
        if isinstance(message_data, bytes):
            message_data = message_data.decode('utf-8')
        message_dict = json.loads(message_data)
        return messages_from_dict([message_dict])[0]

    @property
    def messages(self) -> List[BaseMessage]:
        """Retrieve the most recent messages (up to max_messages)"""
        messages_data = self.redis_client.lrange(self.key, -self.max_messages, -1)
        return [self._deserialize_message(m) for m in messages_data]